from torchvision import datasets
from torchvision import transforms
import torch
import parameters as p
import cv2
import einops
from torch.utils import data
from utils import inject_noise


# transform = transforms.Compose([
#     transforms.ToTensor(),
#     transforms.Resize(p.IMAGE_SIZE),
#     transforms.CenterCrop(p.IMAGE_SIZE),
#     # transforms.Lambda(lambd=lambda t:(t * 2) - 1)
# ])
transform = transforms.Compose([
    transforms.Resize(64, interpolation=transforms.InterpolationMode.BICUBIC),
    transforms.RandomHorizontalFlip(0.4),
    transforms.ToTensor(),  # scale to [0,1]
    # transforms.Lambda(lambda t: (t * 2) - 1)  # scale to [-1,1]
])
# train_dataset = datasets.ImageFolder(p.TRAINING_DATASET,transform=transform)
train_dataset = datasets.MNIST("X:\Machine_Learning\datasets\MNIST_\\",
                         download=False,
                         transform=transform)

def get_dataloader(batch_size=p.BATCH_SIZE,num_workers=4,pin_memory=False,persistent_workers=False):
    global train_dataset
    return data.DataLoader(
        train_dataset,batch_size=batch_size,
        num_workers=num_workers,pin_memory=pin_memory,
        persistent_workers=persistent_workers,drop_last=True)



if __name__ == "__main__":


    # train_dataloader = data.DataLoader(train_dataset,batch_size=p.BATCH_SIZE,num_workers=8)
    img = cv2.imread('L:\\opencv-test\\TEST.jpg')
    # # img2 = cv2.imread("L:\\opencv-test\\sample4.jpg")
    img = transform(img).cuda()
    # # img2 = transform(img2)
    # # img = torch.unsqueeze(img,0)
    # # img2 = torch.unsqueeze(img2,0)
    # print(img.shape)
    # # print(img2.shape)
    
    # # img = torch.concat([img,img2])
    t = torch.tensor([588]).cuda()
    # # print(torch.tensor([400,660]).shape,"#")
    noisy_imgs = inject_noise(img,t=t)
    # print(noisy_imgs.shape)
    
    noisy_imgs = einops.rearrange(noisy_imgs,'1 c h w -> h w c')
    # # img = einops.rearrange(img,'c h w -> h w c')

    cv2.imshow("noisy_image1",noisy_imgs.cpu().numpy())
    # # cv2.imshow("noisy_image2",noisy_imgs[1].numpy())
    cv2.waitKey(0)



    # print(len(train_dataloader))
    # # cv2.namedWindow('sample',cv2.WINDOW_FREERATIO)
    # for data in train_dataloader:
    #     # print(len(data))
    #     # print(data)
    #     img = einops.rearrange(data[0],'1 c h w -> h w c').numpy()
    #     img = cv2.cvtColor(img,cv2.COLOR_RGB2BGR)
    #     cv2.imshow('sample',img)
    #     cv2.waitKey(0)

 
  



